package com.fengye.common.aspectj;

import com.fengye.common.annotation.RateLimit;
import com.fengye.common.constants.AppHttpCodeEnum;
import com.fengye.common.exception.BusinessException;
import com.google.common.collect.Maps;
import com.google.common.util.concurrent.RateLimiter;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.stereotype.Component;

import java.lang.reflect.Method;
import java.util.Map;

/**
 * @author sunlei
 * @version 1.0
 * @date 2024/5/9 10:05:22
 * @description 限流器aop实现（采用令牌桶算法 + Guava 实现）
 */
@Slf4j
@Aspect
@Component
public class RateLimitAspect {
    /**
     * limiter容器
     * 不同的接口，不同的流量控制，通常可以用发方法名作为key，记得要唯一
     * map的key为 Limiter.key
     */
    private final Map<String, RateLimiter> limitMap = Maps.newConcurrentMap();

    @Around("@annotation(com.fengye.common.annotation.RateLimit)")
    public Object around(ProceedingJoinPoint joinPoint) throws Throwable{
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method method = signature.getMethod();
        //拿limit的注解
        RateLimit limit = method.getAnnotation(RateLimit.class);
        if (limit != null) {
            //key作用：不同的接口，不同的流量控制
            String key = limit.key();
            RateLimiter rateLimiter = null;
            //验证缓存是否有命中key，没有就创建limiter
            if (!limitMap.containsKey(key)) {
                // 创建一个指定速率的令牌桶
                rateLimiter = RateLimiter.create(limit.permitsPerSecond());
                limitMap.put(key, rateLimiter);
                log.info("新建了令牌桶={}，容量={}", key, limit.permitsPerSecond());
            }
            rateLimiter = limitMap.get(key);
            // 如果之前的桶的速率跟本次请求设置的速率不一样，则更新桶的速率
            if (rateLimiter.getRate() != limit.permitsPerSecond()) {
                rateLimiter.setRate(limit.permitsPerSecond());
                log.info("令牌桶={}，容量={}，速率={}", key, rateLimiter.getRate(), limit.permitsPerSecond());
            }

            // 判断能否在指定时间内获取到令牌，如果能获取则立即返回true，并消耗一个令牌, 如果超时后依然不能获取令牌，则立即返回 false
            // 扩展：如果想阻塞线程，直到能获取到令牌，则使用 acquire() 方法，这意味着，如果当前没有令牌，线程将会等待，直到有令牌可用或者被中断。
            boolean acquire = rateLimiter.tryAcquire(limit.timeout(), limit.timeunit());
            // 拿不到命令，直接返回异常提示
            if (!acquire) {
                log.warn("令牌桶={}，获取令牌失败，服务降级", key);
                throw new BusinessException(AppHttpCodeEnum.TOO_MANY_REQUESTS);
            }
        }
        return joinPoint.proceed();
    }
}

